#!/usr/bin/env python

import rospy
import sys
import os
import moveit_commander
from moveit_commander import MoveGroupCommander, PlanningSceneInterface
from moveit_msgs.msg import PlanningScene, ObjectColor
from geometry_msgs.msg import PoseStamped, Pose
from math import pi
from copy import deepcopy
import tf
from tf.transformations import *

class MoveItObstaclesDemo:
	def __init__(self):
		moveit_commander.roscpp_initialize(sys.argv)
		rospy.init_node("moveit_obstacles_demo")
		scene = PlanningSceneInterface()
		rospy.sleep(1)
		scene_pub = rospy.Publisher("planning_scene", PlanningScene, queue_size=10)
		arm = moveit_commander.MoveGroupCommander("xarm")
		arm.set_goal_position_tolerance(0.02)
		arm.set_goal_orientation_tolerance(0.03)
		arm.set_max_velocity_scaling_factor(1.0)
		arm.allow_replanning(True)
		arm.set_named_target("Home")
		arm.go()

		table_id = 'table'
		box_id = "box"
		sphere_id = "sphere"
		rospy.loginfo("============Pree 'enter' to add a table to the planning scene...")
		input()
		
		table_size = [1.0, 1.2, 0.01]
		table_pose = PoseStamped()
		table_pose.header.frame_id = "base_link"
		table_pose.pose.position.x = 0.0
		table_pose.pose.position.y = 0.0
		table_pose.pose.position.z = - table_size[2]/2
		table_pose.pose.orientation.w = 1.0

		scene.add_box(table_id, table_pose, table_size)
		if self.wait_for_state_update(table_id, scene):
			rospy.loginfo("The table has been successfully added.")
		else:
			rospy.loginfo("Failed to add the table.")
		rospy.loginfo("============Pree 'enter' to add a sphere to the planning scene...")
		input()
		sphere_radius = 0.12
		sphere_pose = PoseStamped()
		sphere_pose.header.frame_id = "base_link"
		sphere_pose.pose.position.x = 0.3
		sphere_pose.pose.position.y = 0.32
		sphere_pose.pose.position.z = sphere_radius
		sphere_pose.pose.orientation.w = 1.0
		scene.add_sphere(sphere_id, sphere_pose, sphere_radius)
		if self.wait_for_state_update(sphere_id, scene):
			rospy.loginfo("The sphere has been successfully added.")
		else:
			rospy.loginfo("Failed to add the sphere.")
		rospy.loginfo("============Pree 'enter' to add a box to the planning scene...")
		input()
		box_size = [0.25, 0.05, 0.4]
		box_pose = PoseStamped()
		box_pose.header.frame_id = "base_link"
		box_pose.pose.position.x = 0.3
		box_pose.pose.position.y = -0.2
		box_pose.pose.position.z = box_size[2]/2
		box_pose.pose.orientation.w = 1.0
		scene.add_box(box_id, box_pose, box_size)
		if self.wait_for_state_update(box_id, scene):
			rospy.loginfo("The box has been successfully added.")
		else:
			rospy.loginfo("Failed to add the box.")
		rospy.loginfo("============Pree 'enter' to change the color of the box and sphere...")
		input()
		box_color = ObjectColor()
		box_color.id = box_id
		box_color.color.r = 0.8
		box_color.color.g = 0.8
		box_color.color.b = 0.0
		box_color.color.a = 1.0
		sphere_color = ObjectColor()
		sphere_color.id = sphere_id
		sphere_color.color.r = 0.8
		sphere_color.color.g = 0.0
		sphere_color.color.b = 0.9
		sphere_color.color.a = 1.0
		p = PlanningScene()
		p.is_diff = True
		p.object_colors.append(box_color)
		p.object_colors.append(sphere_color)
		scene_pub.publish(p)
		rospy.sleep(1)
		rospy.loginfo("============Pree 'enter' to move the arm...")
		input()
		target_pose = PoseStamped()
		target_pose.header.frame_id = "base_link"
		target_pose.header.stamp = rospy.Time.now()
		target_pose.pose.position.x = 0.25
		target_pose.pose.position.y = -0.4
		target_pose.pose.position.z = 0.25
		quaternion = tf.transformations.quaternion_from_euler(0, 0, -pi/4)
		rospy.loginfo("Quaternion is : ")
		rospy.loginfo(quaternion)
		target_pose.pose.orientation.x = quaternion[0]
		target_pose.pose.orientation.y = quaternion[1]
		target_pose.pose.orientation.z = quaternion[2]
		target_pose.pose.orientation.w = quaternion[3]
		arm.set_start_state_to_current_state()

		end_effector_link = arm.get_end_effector_link()
		arm.set_pose_target(target_pose, end_effector_link)
		arm.go()
		rospy.sleep(1)

		arm.set_named_target("Down")
		arm.go()
		rospy.sleep(1)

		arm.set_named_target("Home")
		arm.go()
		rospy.sleep(1)
		rospy.loginfo("============Pree 'enter' to remove the obstacles and exit the program...")
		input()
		
		scene.remove_world_object(table_id)
		scene.remove_world_object(sphere_id)
		scene.remove_world_object(box_id)
		rospy.sleep(1)

		moveit_commander.roscpp_shutdown()
		moveit_commander.os._exit(0)

	def wait_for_state_update(self, obstacle_name, scene, timeout=4):
		start = rospy.get_time()
		seconds = rospy.get_time()

		while (seconds - start < timeout) and not rospy.is_shutdown():
			is_known = obstacle_name in scene.get_known_object_names()
			if is_known:
				return True

			rospy.sleep(0.1)
			seconds = rospy.get_time()

		return False


if __name__ == "__main__":
	try:
		MoveItObstaclesDemo()
	except rospy.ROSInterruptException:
		pass